import torch
import matplotlib.pyplot as plt


# 准备数据集
x_data = torch.Tensor([[1.0],[2.0],[3.0]]) # 3 * 1 的矩阵
y_data = torch.Tensor([[2.0],[4.0],[6.0]])

# 使用类来设计模型
class LinearModal(torch.nn.Module):
    def __init__(self):
        super(LinearModal,self).__init__()
        self.linear = torch.nn.Linear(1,1)

    def forward(self,x):
        y_pred = self.linear(x)
        return y_pred

modal = LinearModal()

# 构造损失函数和训练器
criterion = torch.nn.MSELoss(reduction='sum')
optimizer = torch.optim.Rprop(modal.parameters(),lr=0.01)

if __name__ == "__main__":

    loss_list = []
    for epoch in range(100):
        y_pred = modal(x_data)
        loss = criterion(y_pred,y_data)
        loss_list.append(loss.item())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    plt.plot(range(100),loss_list)
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.show()

    x_test = torch.Tensor([[4.0]])
    y_test = modal(x_test)
    print('y_test=',y_test)
